{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "cloud_tpu_custom_training.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MhoQ0WE77laV"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "_ckMIh7O7s6D",
        "colab": {}
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jYysdyb-CaWM"
      },
      "source": [
        "# Custom training with TPUs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FbVhjPpzn6BM"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This colab will take you through using [tf.distribute.experimental.TPUStrategy](https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/TPUStrategy). This is a new technique, a part of `tf.distribute.Strategy`, that allows users to easily switch their model to using TPUs. As part of this tutorial, you will create a Keras model and take it through a custom training loop (instead of calling `fit` method).\n",
        "\n",
        "For full information on `DistributionStrategy`, please see the linked documentation above. For the purpose of this Colab, a `DistributionStrategy` is an object that indicates how TF should distribute your model across one or more workers or accelerators. The `TPUStrategy` provides information about your TPU configuration, and we can use it in combination with a Keras model to quickly train that model on the TPU!\n",
        "\n",
        "If you're familiar with Keras models, you may be used to training them with the `model.fit` API. This API is supported using `DistributionStrategy`, but to show the flexibility of using `DistributionStrategy`, we will demonstrate a more customized approach in this notebook, using an explicit training loop."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-3IGpvYjYRBd",
        "colab_type": "text"
      },
      "source": [
        "## Learning objectives\n",
        "\n",
        "In this Colab, you will learn how to\n",
        "* Build a Keras model used to classify image\n",
        "* Set up custom loss, gradient, and training loop functions\n",
        "* Learn to use TPUStrategy so that you can easily port your models from CPU and GPU to TPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zjG11ZawF50B",
        "colab_type": "text"
      },
      "source": [
        "## Instructions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p_e0M_ZpF4CK",
        "colab_type": "text"
      },
      "source": [
        "<h3>  &nbsp;&nbsp;Train on TPU&nbsp;&nbsp; <a href=\"https://cloud.google.com/tpu/\"><img valign=\"middle\" src=\"https://raw.githubusercontent.com/GoogleCloudPlatform/tensorflow-without-a-phd/master/tensorflow-rl-pong/images/tpu-hexagon.png\" width=\"50\"></a></h3>\n",
        "\n",
        "1. On the main menu, click Runtime and select **Change runtime type**. Set \"TPU\" as the hardware accelerator.\n",
        "1. Click Runtime again and select **Runtime > Run All**. You can also run the cells manually with Shift-ENTER. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OYI5KSkYGAge",
        "colab_type": "text"
      },
      "source": [
        "## Imports and hardware and software verification"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "dzLKpmZICaWN",
        "colab": {}
      },
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "\n",
        "# Import TensorFlow\n",
        "import tensorflow as tf\n",
        "\n",
        "# Helper libraries\n",
        "import numpy as np\n",
        "import os\n",
        "\n",
        "assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit > Notebook settings > Hardware accelerator'\n",
        "\n",
        "assert float('.'.join(tf.__version__.split('.')[:2])) >= 1.14, 'Make sure that Tensorflow version is at least 1.14'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5TPCyNk1LjAC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MM6W__qraV55"
      },
      "source": [
        "## Create model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58ff7ew6MK9d",
        "colab_type": "text"
      },
      "source": [
        "Since you will be working with the [MNIST data](https://en.wikipedia.org/wiki/MNIST_database), which is a collection of 70,000 greyscale images representing digits, you want to be using a convolutional neural network to help us with the labeled image data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "7MqDQO0KCaWS",
        "colab": {}
      },
      "source": [
        "def create_model(input_shape):\n",
        "  \"\"\"Creates a simple convolutional neural network model using the Keras API\"\"\"\n",
        "  return tf.keras.Sequential([\n",
        "      tf.keras.layers.Conv2D(28, kernel_size=(3, 3), activation='relu', input_shape=input_shape),\n",
        "      tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
        "      tf.keras.layers.Flatten(),\n",
        "      tf.keras.layers.Dense(128, activation=tf.nn.relu),\n",
        "      tf.keras.layers.Dropout(0.2),\n",
        "      tf.keras.layers.Dense(10, activation=tf.nn.softmax),\n",
        "  ])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4AXoHhrsbdF3"
      },
      "source": [
        "## Loss and gradient"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5mVuLZhbem8d"
      },
      "source": [
        "When using `model.fit`, the gradient function is implemented for you, but for the custom training loop, an explicit loss and gradient function must be defined."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "F2VeZUWUj5S4",
        "colab": {}
      },
      "source": [
        "def loss(model, x, y):\n",
        "  \"\"\"Calculates the loss given an example (x, y)\"\"\"\n",
        "  logits = model(x)\n",
        "  return logits, tf.losses.sparse_softmax_cross_entropy(labels=y, logits=logits)\n",
        "\n",
        "def grad(model, x, y):\n",
        "  \"\"\"Calculates the loss and the gradients given an example (x, y)\"\"\"\n",
        "  logits, loss_value = loss(model, x, y)\n",
        "  return logits, loss_value, tf.gradients(loss_value, model.trainable_variables)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "k53F5I_IiGyI"
      },
      "source": [
        "## Main function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0Qb6nDgxiN_n"
      },
      "source": [
        "Now that we have the model defined, we can load our data, setup the distribution strategy and run our training loop."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "jwJtsCQhHK-E",
        "colab": {}
      },
      "source": [
        "tf.keras.backend.clear_session()\n",
        "\n",
        "resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu=TPU_WORKER)\n",
        "tf.contrib.distribute.initialize_tpu_system(resolver)\n",
        "strategy = tf.contrib.distribute.TPUStrategy(resolver)\n",
        "\n",
        "# Load MNIST training and test data\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "x_train, x_test = x_train / 255.0, x_test / 255.0\n",
        "\n",
        "# All MNIST examples are 28x28 pixel greyscale images (hence the 1\n",
        "# for the number of channels).\n",
        "input_shape = (28, 28, 1)\n",
        "\n",
        "# Only specific data types are supported on the TPU, so it is important to\n",
        "# pay attention to these.\n",
        "# More information:\n",
        "# https://cloud.google.com/tpu/docs/troubleshooting#unsupported_data_type\n",
        "x_train = x_train.reshape(x_train.shape[0], *input_shape).astype(np.float32)\n",
        "x_test = x_test.reshape(x_test.shape[0], *input_shape).astype(np.float32)\n",
        "y_train, y_test = y_train.astype(np.int64), y_test.astype(np.int64)\n",
        "\n",
        "# The batch size must be divisible by the number of workers (8 workers),\n",
        "# so batch sizes of 8, 16, 24, 32, ... are supported.\n",
        "BATCH_SIZE = 32\n",
        "\n",
        "NUM_EPOCHS = 5\n",
        "\n",
        "train_steps_per_epoch = len(x_train) // BATCH_SIZE\n",
        "test_steps_per_epoch = len(x_test) // BATCH_SIZE"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GPrDC8IfOgCT",
        "colab_type": "text"
      },
      "source": [
        "## Start by creating objects within the strategy's scope\n",
        "\n",
        "A key concept to remember: model creation, optimizer creation, etc. must be written in the context of strategy.scope() in order to use TPUStrategy. \n",
        "\n",
        "Also initialize metrics for the train and test sets."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s_suB7CZNw5W",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with strategy.scope():\n",
        "  model = create_model(input_shape)\n",
        "\n",
        "  optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01)\n",
        "\n",
        "  training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)\n",
        "  training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
        "      'training_accuracy', dtype=tf.float32)\n",
        "  test_loss = tf.keras.metrics.Mean('test_loss', dtype=tf.float32)\n",
        "  test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
        "      'test_accuracy', dtype=tf.float32)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d3iLK5ZtO1_R",
        "colab_type": "text"
      },
      "source": [
        "## Define custom train and test steps"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XAF6xfU0N5ID",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with strategy.scope():\n",
        "  def train_step(inputs):\n",
        "    \"\"\"Each training step runs this custom function which calculates\n",
        "    gradients and updates weights.\n",
        "    \"\"\"\n",
        "    x, y = inputs\n",
        "\n",
        "    logits, loss_value, grads = grad(model, x, y)\n",
        "\n",
        "    update_loss = training_loss.update_state(loss_value)\n",
        "    update_accuracy = training_accuracy.update_state(y, logits)\n",
        "\n",
        "    # Show that this is truly a custom training loop\n",
        "    # Multiply all gradients by 2.\n",
        "    grads = grads * 2\n",
        "\n",
        "    update_vars = optimizer.apply_gradients(\n",
        "        zip(grads, model.trainable_variables))\n",
        "\n",
        "    with tf.control_dependencies([update_vars, update_loss, update_accuracy]):\n",
        "      return tf.identity(loss_value)\n",
        "\n",
        "  def test_step(inputs):\n",
        "    \"\"\"Each training step runs this custom function\"\"\"\n",
        "    x, y = inputs\n",
        "\n",
        "    logits, loss_value = loss(model, x, y)\n",
        "\n",
        "    update_loss = test_loss.update_state(loss_value)\n",
        "    update_accuracy = test_accuracy.update_state(y, logits)\n",
        "\n",
        "    with tf.control_dependencies([update_loss, update_accuracy]):\n",
        "      return tf.identity(loss_value)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AhrK1-yEO7Nf",
        "colab_type": "text"
      },
      "source": [
        "## Do the training\n",
        "In order to make the reading a little bit easier, the full training loop calls two helper functions, `run_train()` and `run_test()`. The training runs for 5 epochs and trains to 96% accuracy. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "or5osuheouVU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def run_train():\n",
        "  # Train\n",
        "  session.run(train_iterator_init)\n",
        "  while True:\n",
        "    try:\n",
        "      session.run(dist_train)\n",
        "    except tf.errors.OutOfRangeError:\n",
        "      break\n",
        "  print('Train loss: {:0.4f}\\t Train accuracy: {:0.4f}%'.format(\n",
        "      session.run(training_loss_result),\n",
        "      session.run(training_accuracy_result) * 100))\n",
        "  training_loss.reset_states()\n",
        "  training_accuracy.reset_states()\n",
        "\n",
        "def run_test():\n",
        "  # Test\n",
        "  session.run(test_iterator_init)\n",
        "  while True:\n",
        "    try:\n",
        "      session.run(dist_test)\n",
        "    except tf.errors.OutOfRangeError:\n",
        "      break\n",
        "  print('Test loss: {:0.4f}\\t Test accuracy: {:0.4f}%'.format(\n",
        "      session.run(test_loss_result),\n",
        "      session.run(test_accuracy_result) * 100))\n",
        "  test_loss.reset_states()\n",
        "  test_accuracy.reset_states()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u5LvzAwjN95j",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with strategy.scope():\n",
        "  training_loss_result = training_loss.result()\n",
        "  training_accuracy_result = training_accuracy.result()\n",
        "  test_loss_result = test_loss.result()\n",
        "  test_accuracy_result = test_accuracy.result()\n",
        "  \n",
        "  config = tf.ConfigProto()\n",
        "  config.allow_soft_placement = True\n",
        "  cluster_spec = resolver.cluster_spec()\n",
        "  if cluster_spec:\n",
        "    config.cluster_def.CopyFrom(cluster_spec.as_cluster_def())\n",
        "\n",
        "  print('Starting training...')\n",
        "\n",
        "  # Do all the computations inside a Session (as opposed to doing eager mode)\n",
        "  with tf.Session(target=resolver.master(), config=config) as session:\n",
        "    all_variables = (\n",
        "        tf.global_variables() + training_loss.variables +\n",
        "        training_accuracy.variables + test_loss.variables +\n",
        "        test_accuracy.variables)\n",
        "    \n",
        "    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(BATCH_SIZE, drop_remainder=True)\n",
        "    train_iterator = strategy.make_dataset_iterator(train_dataset)\n",
        "\n",
        "    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(BATCH_SIZE, drop_remainder=True)\n",
        "    test_iterator = strategy.make_dataset_iterator(train_dataset)\n",
        "    \n",
        "    train_iterator_init = train_iterator.initialize()\n",
        "    test_iterator_init = test_iterator.initialize()\n",
        "\n",
        "    session.run([v.initializer for v in all_variables])\n",
        "    \n",
        "    dist_train = strategy.experimental_run(train_step, train_iterator).values\n",
        "    dist_test = strategy.experimental_run(test_step, test_iterator).values\n",
        "\n",
        "    # Custom training loop\n",
        "    for epoch in range(0, NUM_EPOCHS):\n",
        "      print('Starting epoch {}'.format(epoch))\n",
        "\n",
        "      run_train()\n",
        "\n",
        "      run_test()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UkWkS928YDA-",
        "colab_type": "text"
      },
      "source": [
        "## What's next\n",
        "\n",
        "* Learn about [Cloud TPUs](https://cloud.google.com/tpu/docs) that Google designed and optimized specifically to speed up and scale up ML workloads for training and inference and to enable ML engineers and researchers to iterate more quickly.\n",
        "* Explore the range of [Cloud TPU tutorials and Colabs](https://cloud.google.com/tpu/docs/tutorials) to find other examples that can be used when implementing your ML project.\n",
        "\n",
        "On Google Cloud Platform, in addition to GPUs and TPUs available on pre-configured [deep learning VMs](https://cloud.google.com/deep-learning-vm/),  you will find [AutoML](https://cloud.google.com/automl/)*(beta)* for training custom models without writing code and [Cloud ML Engine](https://cloud.google.com/ml-engine/docs/) which will allows you to run parallel trainings and hyperparameter tuning of your custom models on powerful distributed hardware.\n"
      ]
    }
  ]
}